library(tidyverse)
library(ggplot2)
library(ggrepel)
library(readr)
library(lme4)
library(interactions)
library(patchwork)
library(ggridges)
library(jtools)
library(modelsummary)
library(colorspace)
library(interplot)
library(ggpubr)
library(margins)
library(estimatr)
library(sandwich)
library(patchwork)
library(stargazer)


# Load the data
data <- read_dta("Study_1.dta")

# Create treatment variable
data <- data %>%
  mutate(treatment = case_when(
    trustExperimentSplit == 1 ~ 0,
    trustExperimentSplit == 2 ~ 1
  )) %>%
  mutate(treatment = factor(treatment, labels = c("Control", "Treatment")))

# Create trust variables
data <- data %>%
  mutate(
    trustMP = coalesce(trustExperimentGroup1MPRe, trustExperimentGroup2MPRe),
    trustMSP = coalesce(trustExperimentGroup1MSPRe, trustExperimentGroup2MSPRe),
    trustScotMin = coalesce(trustExperimentGroup1ScotMinRe, trustExperimentGroup2ScotMinRe),
    trustUKMin = coalesce(trustExperimentGroup1UKMinRe, trustExperimentGroup2UKMinRe),
    trustCivWest = coalesce(trustExperimentGroup1CivWestRe, trustExperimentGroup2CivWestRe),
    trustCivScot = coalesce(trustExperimentGroup1CivScotRe, trustExperimentGroup2CivScotRe)
  )

# Create binary independence variable
data <- data %>% 
  mutate(indy01 = case_when(
    voteIntIndy == 2 ~ 0,
    voteIntIndy == 1 ~ 1,
    TRUE ~ NA_real_
  ))

# Create subsets for independence supporters and non-supporters
indy_supporters <- subset(data, indy01 == 1)
union_supporters <- subset(data, indy01 == 0)

# Function to run models and create plots for each outcome
create_outcome_plot <- function(outcome_var, data_indy, data_union) {
  # Determine overall range for y-axis
  y_min <- min(c(data_indy[[outcome_var]], data_union[[outcome_var]]), na.rm = TRUE)
  y_max <- max(c(data_indy[[outcome_var]], data_union[[outcome_var]]), na.rm = TRUE)
  
  y_range <- y_max - y_min
  buffer <- 0.10 * y_range
  
  y_limits <- c(max(0, y_min - buffer), min(5, y_max + buffer))
  
  # Models
  model_indy <- lm(as.formula(paste(outcome_var, "~ treatment")), data = data_indy)
  model_union <- lm(as.formula(paste(outcome_var, "~ treatment")), data = data_union)
  
  # Predictions
  data_indy$predicted <- predict(model_indy, data_indy)
  data_union$predicted <- predict(model_union, data_union)
  
  # Subsets for treatment and control
  treat_indy <- subset(data_indy, treatment == "Treatment")
  control_indy <- subset(data_indy, treatment == "Control")
  treat_union <- subset(data_union, treatment == "Treatment")
  control_union <- subset(data_union, treatment == "Control")
  
  # Calculate ATE and p-values
  ate_indy <- coef(model_indy)["treatmentTreatment"]
  ate_union <- coef(model_union)["treatmentTreatment"]
  p_value_indy <- summary(model_indy)$coefficients["treatmentTreatment", "Pr(>|t|)"]
  p_value_union <- summary(model_union)$coefficients["treatmentTreatment", "Pr(>|t|)"]
  
  # Function to add asterisks based on p-value
  add_asterisks <- function(p_value) {
    if (p_value < 0.01) return("***")
    else if (p_value < 0.05) return("**")
    else if (p_value < 0.1) return("*")
    else return("")
  }
  
  # Create plots
  plot_indy <- effect_plot(model = model_indy, pred = treatment, robust = TRUE,
                           cat.geom = "point", cat.interval.geom = "linerange",
                           colors = "black", cat.pred.point.size = 3, int.width = .90) +
    labs(title = paste("Independence Supporters (N =", nrow(data_indy), ")")) +
    ylab(paste("Trust in", gsub("trust", "", outcome_var))) +
    xlab("") +
    scale_x_discrete(labels = c("Control" = "Control", "Treatment" = "Treatment")) +
    scale_y_continuous(limits = y_limits) +
    geom_jitter(data = treat_indy, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, size = 3, width = .35, alpha = .15, shape = 20,
                pch = 21, color = "#11CE9E") +
    geom_jitter(data = control_indy, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 17,
                pch = 21, size = 3, color = "#CE1141FF") +  
    geom_bracket(xmin = "Control", xmax = "Treatment",
                 y.position = y_limits[2] * 0.95,
                 label = sprintf("ATE=%.2f%s", ate_indy, add_asterisks(p_value_indy)),
                 tip.length = 0.02,
                 color = "black") +
    theme_minimal(base_size = 14) +
    theme(axis.title.y = element_text(face = "bold"),
          axis.text.x = element_text(face = "bold"))
  
  plot_union <- effect_plot(model = model_union, pred = treatment, robust = TRUE,
                            cat.geom = "point", cat.interval.geom = "linerange",
                            colors = "black", cat.pred.point.size = 3, int.width = .90) +
    labs(title = paste("Union Supporters (N =", nrow(data_union), ")")) +
    ylab("") +
    xlab("") +
    scale_x_discrete(labels = c("Control" = "Control", "Treatment" = "Treatment")) +
    scale_y_continuous(limits = y_limits) +
    geom_jitter(data = treat_union, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 20,
                pch = 21, size = 3, color = "#11CE9E") +
    geom_jitter(data = control_union, aes(x = treatment, y = .data[[outcome_var]]),
                height = .25, width = .35, alpha = .15, shape = 17,
                pch = 21, size = 3, color = "#CE1141FF") +
    geom_bracket(xmin = "Control", xmax = "Treatment",
                 y.position = y_limits[2] * 0.95,
                 label = sprintf("ATE=%.2f%s", ate_union, add_asterisks(p_value_union)),
                 tip.length = 0.02,
                 color = "black") +
    theme_minimal(base_size = 14) +
    theme(axis.title.y = element_blank(),
          axis.text.y = element_blank(),
          axis.text.x = element_text(face = "bold"))
  
  # Combine plots and return
  plot_indy + plot_union +
    plot_annotation(title = paste('Effect of treatment on trust in', gsub("trust", "", outcome_var)),
                    theme = theme(plot.title = element_text(size = 14, face = "bold")))
}

# Create plots for each outcome
outcomes <- c("trustMP", "trustMSP", "trustUKMin", "trustScotMin", "trustCivWest", "trustCivScot")
plots <- map(outcomes, ~create_outcome_plot(.x, indy_supporters, union_supporters))

# Combine all plots into one
combined_plot <- wrap_plots(plots, ncol = 2) +
  plot_annotation(
    title = 'Effects of treatment on trust in various entities',
    caption = "Treatment group outcome statistically distinct at p<.1(*), p<0.05(**), & p<0.01(***)",
    theme = theme(plot.title = element_text(size = 20, face = "bold"))
  )

# Save the combined plot
ggsave("Figure_3.png", combined_plot, width = 20, height = 30, dpi = 300)

# Create models
models <- list()
for (outcome in outcomes) {
  models[[outcome]] <- list(
    all = lm(as.formula(paste(outcome, "~ treatment")), data = data),
    interaction = lm(as.formula(paste(outcome, "~ treatment * indy01")), data = data),
    indy = lm(as.formula(paste(outcome, "~ treatment")), data = indy_supporters),
    union = lm(as.formula(paste(outcome, "~ treatment")), data = union_supporters)
  )
}

# Create and save tables
for (outcome in outcomes) {
  stargazer(models[[outcome]], 
            type = "latex",
            title = paste("Models for", outcome),
            column.labels = c("All", "Interaction", "Independence", "Union"),
            out = paste0("Table_", outcome, ".tex"))
}
})